// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::ffi::CStr;

use indexmap::IndexMap;
use once_cell::sync::Lazy;
use tracing::debug;

use super::kninfo_gen::{build_func_map, KNINFO_LIST};

static FUNC_MAP: Lazy<IndexMap<&'static str, usize>> =
    Lazy::new(|| build_func_map(KNINFO_LIST.len()));

#[inline]
fn get_func_index(func_name: &str) -> Option<usize> {
    FUNC_MAP.get(func_name).copied()
}

#[inline]
fn get_kninfo_by_index(index: usize) -> Option<(&'static str, &'static [usize])> {
    KNINFO_LIST.get(index).copied()
}

#[inline]
pub fn get_kninfo(func_name: &CStr) -> Option<(&'static str, &'static [usize])> {
    let func_name = func_name.to_str().expect("Unexpected C String");
    debug!("[kninfo] get_kninfo({:?})", func_name);

    let index = get_func_index(func_name)?;
    debug!("[kninfo] get_kninfo, index={}", index);

    get_kninfo_by_index(index)
}

#[cfg(test)]
mod tests {
    use crate::utils::kninfo::*;

    #[test]
    fn test_kn_info_map() {
        let key = "_ZN2at6native18elementwise_kernelILi128ELi2EZNS0_22gpu_kernel_impl_nocastINS0_15CUDAFunctor_addIfEEEEvRNS_18TensorIteratorBaseERKT_EUliE_EEviT1_";
        let index = get_func_index(key).expect("get_func_index_failed");
        println!("index={}", index);
        let _ = get_kninfo_by_index(index as usize).expect("get_kninfo_by_index");
    }
}
